package vulnerability

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/go-jet/jet/v2/postgres"
	"github.com/google/uuid"
	"github.com/rs/zerolog/log"
	"go.uber.org/fx"

	"github.com/lunasec-io/lunasec/lunatrace/bsl/ingest-worker/pkg/ml"
	"github.com/lunasec-io/lunasec/lunatrace/gogen/sqlgen/lunatrace/vulnerability/model"
	"github.com/lunasec-io/lunasec/lunatrace/gogen/sqlgen/lunatrace/vulnerability/table"
)

// TODO (cthompson) this is similar to metadata/processor and should be refactored to be more generic
// if specific logic for each implementation is needed, use options
type Processor interface {
	GenerateEmbeddingsForReferences(vulnID string) error
}

type Params struct {
	fx.In

	DB *sql.DB
	ML ml.Service
}

type processor struct {
	deps Params
}

func (p *processor) GenerateEmbeddingsForReferences(vulnID string) error {
	rc := table.ReferenceContent
	r := table.Reference
	v := table.Vulnerability

	getReferenceContentStmt := rc.SELECT(rc.AllColumns, r.AllColumns, v.ID).
		FROM(
			rc.LEFT_JOIN(r, r.ID.EQ(rc.ReferenceID)).
				LEFT_JOIN(v, v.ID.EQ(r.VulnerabilityID)),
		)
	if vulnID != "" {
		getReferenceContentStmt = getReferenceContentStmt.
			WHERE(v.SourceID.EQ(postgres.String(vulnID)))
	}

	rows, err := getReferenceContentStmt.Rows(context.Background(), p.deps.DB)
	if err != nil {
		log.Error().Err(err).Msg("failed to get reference content")
		return err
	}

	for rows.Next() {
		var ref struct {
			model.Reference
			model.ReferenceContent
			model.Vulnerability
		}

		err = rows.Scan(&ref)
		if err != nil {
			log.Error().Err(err).Msg("failed to scan reference content")
			return err
		}

		// TODO (cthompson) this is a hack to get the vuln id and title into the reference content
		ref.NormalizedContent = fmt.Sprintf("%s %s %s", ref.Vulnerability.SourceID, ref.Title, ref.NormalizedContent)

		normalizedRef := ml.ReferenceContent{
			ID:                  ref.ReferenceContent.ID,
			URL:                 ref.URL,
			Content:             ref.Content,
			NormalizedContent:   ref.NormalizedContent,
			ContentType:         ref.ContentType,
			LastSuccessfulFetch: ref.LastSuccessfulFetch,
		}

		var re = table.ReferenceEmbedding
		insertRefEmbedding := func(id uuid.UUID, contentHash, content, embedding string) error {
			newRefEmb := model.ReferenceEmbedding{
				ReferenceContentID: id,
				ContentHash:        contentHash,
				Content:            contentHash,
				Embedding:          embedding,
			}

			insertStmt := re.INSERT(
				re.ReferenceContentID, re.ContentHash, re.Content, re.Embedding,
			).MODEL(newRefEmb)

			_, err = insertStmt.Exec(p.deps.DB)
			return err
		}

		refEmbeddingExists := func(contentHash string) (string, bool) {
			getExistingRefEmb := re.SELECT(
				re.Embedding,
			).WHERE(re.ContentHash.EQ(postgres.String(contentHash)))

			var refEmb model.ReferenceEmbedding
			err := getExistingRefEmb.Query(p.deps.DB, &refEmb)
			return refEmb.Embedding, err == nil
		}

		err = p.deps.ML.GenerateEmbeddingForRef(&normalizedRef, refEmbeddingExists, insertRefEmbedding)
		if err != nil {
			log.Error().
				Err(err).
				Str("url", ref.URL).
				Msg("failed to generate embedding for reference")
			return err
		}
	}
	return nil
}

func NewProcessor(deps Params) Processor {
	return &processor{
		deps: deps,
	}
}
